﻿
#-----------------------------------------------------------------
# Imports
import bpy
import math
import re
import mathutils

from myxml import *
from tags import *
    
#---------------------------------------------

# Obj should be an armature object
def write(obj, filename):
    
    if(obj.type != 'ARMATURE'):
        print('Anim: selected object not an armature')
        return

    file = open(filename, 'w')

    def w(text) :
        file.write(text)
    
    write_actions(obj, w)
    
    file.close()
    
    print("Object '" + obj.name + "' : anims exported")


# parameter w is a function that you call with string-data
# to write to a file
def write_actions(obj, w):   
    
    # Save state
    orig_scene_frame = bpy.context.scene.frame_current
    orig_scene_preview_start = bpy.context.scene.frame_preview_start
    orig_scene_preview_end = bpy.context.scene.frame_preview_end
    orig_action = obj.animation_data.action
    
    #- - - - - - - - - - -
    
    arm = obj.data # This is the actual armature data
    
    arm_bone_names = [bone.name for bone in obj.pose.bones]
    
    # custom bone sets (allows overlapping "bone groups", which Blender
    # does not allow).
    tags = get_tags_dict(obj)
    print(tags)     
    
    #- - - - - - - - - - -
    
    # First we count the number of actions that will be exported
    num_actions = 0
    for action in bpy.data.actions:        
        action_tags = get_action_tags(action)                
        # Check if the action speaks of any bones in the armature
        if len(action_tags.intersection(tags)) != 0:
            num_actions += 1 # it does, so this action will be exported
            
    #- - - - - - - - - - -
    
    w(otag('actions', [('count', num_actions)]))
    
    for action in bpy.data.actions:
        
        action_tags = get_action_tags(action)
                
        # Check if the action speaks of any bones in the armature
        if len(action_tags.intersection(tags)) == 0:
            continue # it does not, so skip this action
 
        # Otherwise we sample key frames from this action,
        # for the bones that it mentions. These are the bones
        # that should be exported:        
        export_bone_names = set()
        for tag in action_tags:            
            tagged_bones = tags.get(tag)
            if(tagged_bones is None):
                continue # the action mentions an unused tag            
            for bone_name in tagged_bones:
                export_bone_names.add(bone_name)
        
        act_start = int(action.frame_range[0])
        act_end = int(action.frame_range[1])
                
        w(otag('action', [('name', action.name), \
                          ('frame_start', act_start), \
                          ('frame_end', act_end)]))        

        bpy.context.scene.frame_preview_start = act_start 
        bpy.context.scene.frame_preview_end = act_end
         
        obj.animation_data.action = action
        
        # Just to be sure: clear any transformations on the pose bones. This
        # is because some bones are not directly keyed in the action, nor influenced
        # by any keyed bones, and hence might retain some transformed state from
        # a previous action. 
        set_to_binding_pose(obj)
        
        bpy.context.scene.update()
 
        for bone_name in export_bone_names:
            
            w(otag('bone', [('name', bone_name)]))
            
            p_bone = obj.pose.bones[bone_name] # pose bone
            a_bone = obj.data.bones[bone_name] # armature bone (for binding pose)
            
            # binding pose: get relative transform of bone with respect to parent
            m1 = mathutils.Matrix()            
            if a_bone.parent is not None:
                m1 = a_bone.parent.matrix_local.inverted()
            m1 = m1 * a_bone.matrix_local
            
            # Go over the frames (export_fbx, line 2679)
            # Note: we do plus 1 in for the range-stop, because 'stop' is exclusive            
            for frame in range(act_start, act_end+1):
                
                bpy.context.scene.frame_set(frame)
                
                # animated pose: get relative transform of bone with respect to parent                
                m2 = mathutils.Matrix()
                if p_bone.parent is not None:
                    m2 = p_bone.parent.matrix.inverted()
                m2 = m2 * p_bone.matrix
                
                # Get relative transform of animated pose with respect to binding pose (relative to parent)
                m3 = m1.inverted() * m2      
                
                # Make the bone-space use the Z-axis as the "up" axis 
                # (consistent with the scene) --> OOPS, not needed                
                #corr1 = mathutils.Matrix.Rotation(math.radians(-90.0),4,'X')
                #corr2 = mathutils.Matrix.Rotation(math.radians(90.0),4,'X')    
                # Read from right to left
                #m3 = corr2 * m3 * corr1          
                
                scale = m3.to_scale() # float array [3]
                location = m3.to_translation() # float array [3]
                rotation = m3.to_quaternion() # float array [4], w,x,y,z                             
                
                w(otag('f', [('nr', frame)]))
                                
                # String = list of characters, so zip is possible
                w(octag('loc', zip('xyz', location)))
                w(octag('rot', zip('wxyz', rotation)))
                w(octag('scale', zip('xyz', scale)))
                
                w(ctag('f'))
                
            w(ctag("bone"))
        
        w(ctag('action'))   
        
    w(ctag('actions'))
    
    #- - - - - - - - - - -
    # Cleanup
            
    set_to_binding_pose(obj) # Todo: for user comfort, perhaps save the pose at the beginning and then restore it here
    obj.animation_data.action = orig_action
                
    # restore scene
    bpy.context.scene.frame_current = orig_scene_frame 
    bpy.context.scene.frame_preview_start = orig_scene_preview_start 
    bpy.context.scene.frame_preview_end = orig_scene_preview_end 
    bpy.context.scene.update()



#-----------------------------------------------------------------

# Get a set of all tags mentioned in the action name
def get_action_tags(action):
    
    print("Action '" + action.name +"', tags: ")
        
    p = re.compile(r'([^\|]+)\|.*')
    m = p.match(action.name)
    
    if m is None:
        # No tags are given
        print("\t none")
        return set()    
    
    # Otherwise, collect tags
    tags = set()
    
    group = m.group(1)
    for tag in group.split(","):
        tag = tag.strip()
        print("\t " + tag)
        tags.add(tag)
    
    return tags

#-----------------------------------------------------------------
    
def set_to_binding_pose(obj):
    for bone in obj.pose.bones:
        bone.rotation_quaternion.identity()
        bone.scale.xyz = 1
        bone.location.xyz = 0
        
#-----------------------------------------------------------------
    
# Copied from export_fbx (previously called action_bone_names)
# This function gets the bone names that are mentioned in an action.
def get_action_bone_names(obj, action):
    from bpy.types import PoseBone

    names = set()
    path_resolve = obj.path_resolve

    for fcu in action.fcurves:
        try:
            prop = path_resolve(fcu.data_path, False)
        except:
            prop = None

        if prop is not None:
            data = prop.data
            if isinstance(data, PoseBone):
                names.add(data.name)

    return names

